# coding:utf-8

"""Ner-bert-pure predict

Author:
    name: reeseimk
    email: reeseimk@163.com

Homepage: https://gitee.com/reeseimk/mindspore_bert
"""

chunk_dict = {"地名": ["B-Loc", "I-Loc"]}
label_map = {"B-Loc": 0, "I-Loc": 1, "O": 2}
text = "清华大学坐落于首都北京"
labels = []

# text = "重庆是一个魔幻城市"
# labels = [[0,2,"地名"]]

def get_ner_label(text, labels, max_len):
    res_label = [label_map["O"]] * len(text)
    if labels:
        
        for label in labels:
            b, i = chunk_dict[label[2]]
            b_ids, i_ids = label_map[b], label_map[i]
            # print(b_ids, i_ids)
            res_label[label[0]] = b_ids
            i_len = label[1] - 1
            if i_len:
                for i_count in range(1, i_len+1):
                    res_label[label[0]+i_count] = i_ids
    res_label = [label_map["O"]] + res_label + [label_map["O"]] + [label_map["O"]] * (max_len - len(text))
    res_label = res_label[:max_len]
    print(res_label)
    return res_label

if __name__ == "__main__":
    get_ner_label(text, labels, 5)